import os
import io
import torch
from PIL import Image
import streamlit as st
from diffusers import QwenImageEditPipeline

# ---------------------------
# 初始化模型（只加载一次，带加载动画）
# ---------------------------
@st.cache_resource
def load_pipeline():
    model_path = "/home/use/qwen-image-edit-4bit"
    pipe = QwenImageEditPipeline.from_pretrained(model_path, torch_dtype=torch.bfloat16)
    pipe.set_progress_bar_config(disable=None)
    pipe.to("cuda")   # 如果显存不够可换成 pipe.enable_model_cpu_offload()
    return pipe

with st.spinner("🚀 正在加载模型，请稍候..."):
    pipeline = load_pipeline()

# ---------------------------
# Streamlit 页面配置
# ---------------------------
st.set_page_config(page_title="Qwen Image Editor", page_icon="🎨", layout="wide")
st.title("🎨 Qwen Image Editor")

# ---------------------------
# 多会话管理
# ---------------------------
if "conversations" not in st.session_state:
    st.session_state.conversations = {}  # {id: {"history": [], "current_image": None, "name": str}}
if "current_conv" not in st.session_state:
    st.session_state.current_conv = None

# ---------------------------
# 左侧栏：会话管理
# ---------------------------
with st.sidebar:
    st.header("💬 会话管理")

    # 新建对话
    if st.button("🆕 新建对话"):
        new_id = f"conv_{len(st.session_state.conversations) + 1}"
        st.session_state.conversations[new_id] = {
            "history": [],
            "current_image": None,
            "name": f"对话 {len(st.session_state.conversations) + 1}"
        }
        st.session_state.current_conv = new_id
        st.rerun()

    # 会话选择
    if st.session_state.conversations:
        names = [conv["name"] for conv in st.session_state.conversations.values()]
        ids = list(st.session_state.conversations.keys())
        current_index = ids.index(st.session_state.current_conv) if st.session_state.current_conv else 0

        selected_name = st.selectbox("选择会话", names, index=current_index)
        selected_conv = ids[names.index(selected_name)]
        st.session_state.current_conv = selected_conv

        # 会话重命名
        new_name = st.text_input("✏️ 重命名当前会话", st.session_state.conversations[selected_conv]["name"])
        st.session_state.conversations[selected_conv]["name"] = new_name

    else:
        st.info("请先新建一个对话。")

# ---------------------------
# 当前会话对象
# ---------------------------
if st.session_state.current_conv:
    conv = st.session_state.conversations[st.session_state.current_conv]

    # ---------------------------
    # 上传初始图片（中文化 + 自定义按钮样式）
    # ---------------------------
    st.subheader("📤 上传图片")

    # 注入 CSS：彻底隐藏英文 Browse files，并只显示中文
    custom_upload_style = """
    <style>
    /* 隐藏 file_uploader 默认的所有英文提示 */
    [data-testid="stFileUploaderDropzone"] * {
        color: transparent !important;
        font-size: 0 !important;
        visibility: hidden !important;
        height: 0 !important;
        min-height: 0 !important;
        margin: 0 !important;
        padding: 0 !important;
    }
    /* 简约上传区域样式 */
    [data-testid="stFileUploaderDropzone"] {
        position: relative;
        border: none !important;
        border-radius: 10px !important;
        padding: 28px 0 28px 0;
        text-align: center;
        background-color: #f5f5f5 !important;
        cursor: pointer;
        transition: background 0.2s;
        font-size: 16px;
        font-weight: 500;
        min-height: 60px;
    }
    [data-testid="stFileUploaderDropzone"]:hover {
        background-color: #ececec !important;
    }
    /* 添加中文自定义文本 */
    [data-testid="stFileUploaderDropzone"]::after {
        content: "点击上传图片";
        color: #666;
        font-size: 16px;
        font-weight: 500;
        visibility: visible !important;
        display: block;
        position: absolute;
        left: 0; right: 0; top: 50%; transform: translateY(-50%);
        pointer-events: none;
    }
    </style>
    """
    st.markdown(custom_upload_style, unsafe_allow_html=True)

    upload_col, preview_col = st.columns([1, 2])

    with upload_col:
        uploaded_file = st.file_uploader(
            "上传图片",   # 提供非空 label
            type=["png", "jpg", "jpeg"],
            key=f"upload_{st.session_state.current_conv}",
            label_visibility="collapsed"  # 隐藏 label
        )
        if uploaded_file is not None and conv["current_image"] is None:
            conv["current_image"] = Image.open(uploaded_file).convert("RGB")
            conv["history"] = []  # 重置对话

    with preview_col:
        if conv["current_image"] is not None:
            st.image(
                conv["current_image"],
                caption="当前会话图片",
                width="stretch",  # 替换 use_container_width
                output_format="PNG"
            )

    # ---------------------------
    # 显示历史（下载按钮唯一 key）
    # ---------------------------
    for i, (prompt, img) in enumerate(conv["history"]):
        with st.chat_message("user"):
            st.write(prompt)
        with st.chat_message("assistant"):
            st.image(img, caption=f"结果 {i+1}", width=400)
            buf = io.BytesIO()
            img.save(buf, format="PNG")
            st.download_button(
                label="⬇️ 下载图片",
                data=buf.getvalue(),
                file_name=f"{conv['name']}_result_{i+1}.png",
                mime="image/png",
                key=f"download_{st.session_state.current_conv}_{i+1}"
            )

    # ---------------------------
    # 用户输入 prompt
    # ---------------------------
    if conv["current_image"] is not None:
        prompt = st.chat_input("请输入图片编辑指令(越详细效果越好)")
        if prompt:
            with st.chat_message("user"):
                st.write(prompt)

            with st.chat_message("assistant"):
                with st.spinner("🎨 正在生成，请稍候..."):
                    inputs = {
                        "image": conv["current_image"],
                        "prompt": prompt,
                        "generator": torch.manual_seed(0),
                        "true_cfg_scale": 4.0,
                        "negative_prompt": " ",
                        "num_inference_steps": 20,
                    }

                    with torch.inference_mode():
                        output = pipeline(**inputs)

                    result_image = output.images[0]

                    # 更新历史
                    conv["history"].append((prompt, result_image))
                    conv["current_image"] = result_image  # 下一轮编辑的起点

                    # 显示结果
                    st.image(result_image, caption=f"结果 {len(conv['history'])}", width=400)

                    # 添加下载按钮（唯一 key）
                    buf = io.BytesIO()
                    result_image.save(buf, format="PNG")
                    st.download_button(
                        label="⬇️ 下载图片",
                        data=buf.getvalue(),
                        file_name=f"{conv['name']}_result_{len(conv['history'])}.png",
                        mime="image/png",
                        key=f"download_{st.session_state.current_conv}_{len(conv['history'])}"
                    )
    else:
        st.info("请先上传一张图片。")
